{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.core.display import display, HTML\n",
    "display(HTML(\"<style>.container { width:100% !important; }</style>\"))\n",
    "\n",
    "import warnings\n",
    "warnings.filterwarnings('ignore')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from astropy.io import fits\n",
    "from astronet.preprocess import generate_input_records\n",
    "from astronet.preprocess import preprocess\n",
    "from light_curve_util import keplersplinev2\n",
    "from light_curve_util import tess_io\n",
    "from matplotlib import pyplot as plt\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "\n",
    "\n",
    "tce_tables = {}\n",
    "\n",
    "tce_files = [\n",
    "    '/mnt/tess/astronet/tces-toi.csv',\n",
    "#     '/mnt/tess/astronet/tces-toi-bls-vs-tev.csv',\n",
    "    '/mnt/tess/astronet/tces-s33_cam1_sample.csv',\n",
    "    '/mnt/tess/astronet/tces-s33_cam2ccd14_sample.csv',\n",
    "    '/mnt/tess/astronet/tces-v6-train.csv',\n",
    "    '/mnt/tess/astronet/tces-v6-val.csv',\n",
    "    '/mnt/tess/astronet/tces-v6-test.csv',\n",
    "]\n",
    "\n",
    "\n",
    "def get_tce(tic, item_no):\n",
    "    for f in tce_files:\n",
    "        if f not in tce_tables:\n",
    "            tce_tables[f] = pd.read_csv(f, header=0)\n",
    "        tce_table = tce_tables[f]\n",
    "        tce = tce_table[tce_table.tic_id == tic]\n",
    "        if not len(tce):\n",
    "            continue\n",
    "        if len(tce) > 1:\n",
    "            tce = tce[tce.index == tce.index.values[item_no]]\n",
    "            if 'Source' in tce:\n",
    "                print('Source:', 'BLS' if tce.Source.values.item() == 2 else 'TEV')\n",
    "        return tce, f\n",
    "    raise ValueError(f'no TCE data for {tic}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_all(tic,\n",
    "             large=False,\n",
    "             dpi=100,\n",
    "             bkspace=None,\n",
    "             abs_xlim=None,\n",
    "             abs_offset=0,\n",
    "             abs_ylim=None,\n",
    "             rpt_only=False,\n",
    "             item_no=0):\n",
    "    tess_data_dir = f'/mnt/tess/lc'\n",
    "    reports_dir = f'/mnt/tess/rpt/png'\n",
    "\n",
    "    fsize = (16, 9)\n",
    "    \n",
    "    tce, tces_file = get_tce(tic, item_no)\n",
    "    period = tce.Period.values.item()\n",
    "    epoc = tce.Epoc.values.item()\n",
    "    duration = tce.Duration.values.item()\n",
    "\n",
    "    print(f'Epoc: {epoc}\\nPeriod: {period}\\nDuration: {duration}')\n",
    "    \n",
    "    generate_input_records.FLAGS = generate_input_records.parser.parse_args([\n",
    "      '--tess_data_dir', '/mnt/tess/lc',\n",
    "      '--output_dir', '/dev/null',\n",
    "      '--input_tce_csv_file', tces_file,\n",
    "    ])\n",
    "\n",
    "\n",
    "    row = list(tce.iterrows())[0][1]\n",
    "\n",
    "    ex = generate_input_records._process_tce(row, bkspace)\n",
    "    return ex\n",
    "\n",
    "def plot_normal(ex):\n",
    "    fig = plt.figure()\n",
    "    ax = fig.add_subplot(211)\n",
    "    ax.set_xticklabels([])\n",
    "    plt.rcParams['xtick.direction'] = 'in'\n",
    "    plt.rcParams['ytick.direction'] = 'in'\n",
    "    plt.plot(ex.features.feature['local_view'].float_list.value, linestyle='', marker='.')\n",
    "    \n",
    "\n",
    "ex = plot_all(25155310)\n",
    "plot_normal(ex)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_img(ex):\n",
    "    fig = plt.figure(figsize=(6, 2))\n",
    "    ax = fig.add_subplot(111)\n",
    "    ax.axes.get_xaxis().set_visible(False)\n",
    "    ax.axes.get_yaxis().set_visible(False)\n",
    "    ax.spines['bottom'].set_color('white')\n",
    "    ax.spines['right'].set_color('white')\n",
    "    ax.spines['top'].set_color('white')\n",
    "    ax.spines['left'].set_color('white')\n",
    "    plt.imshow(\n",
    "        np.array(ex.features.feature['local_view'].float_list.value).reshape(1, 61),\n",
    "        cmap=plt.get_cmap('bwr'),\n",
    "        aspect='auto',\n",
    "        interpolation='none',\n",
    "    )\n",
    "    \n",
    "plot_img(ex)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "def plot_all(tic,\n",
    "             large=False,\n",
    "             dpi=100,\n",
    "             bkspace=None,\n",
    "             abs_xlim=None,\n",
    "             abs_offset=0,\n",
    "             abs_ylim=None,\n",
    "             rpt_only=False,\n",
    "             item_no=0):\n",
    "    tess_data_dir = f'/mnt/tess/lc'\n",
    "    reports_dir = f'/mnt/tess/rpt/png'\n",
    "\n",
    "    if rpt_only:\n",
    "        plt.figure(figsize=(12, 8), dpi=dpi)\n",
    "        plt.axis('off')\n",
    "        im = plt.imread(f'{reports_dir}/TIC{tic}.png')\n",
    "        _ = plt.imshow(im)\n",
    "        plt.show()\n",
    "        print(get_tce(tic, item_no)[0])\n",
    "        return\n",
    "        \n",
    "    b = plt.get_cmap('tab20')(0)\n",
    "    b2 = plt.get_cmap('tab20')(1)\n",
    "    o = plt.get_cmap('tab20')(2)\n",
    "    o2 = plt.get_cmap('tab20')(3)\n",
    "    g = plt.get_cmap('tab20')(4)\n",
    "    r = plt.get_cmap('tab20')(6)\n",
    "    n = plt.get_cmap('tab20')(8)\n",
    "    k = plt.get_cmap('tab20')(14)\n",
    "\n",
    "    plotrows = 1\n",
    "    plotcols = 3\n",
    "    \n",
    "    if large:\n",
    "        fsize = (16, 9)\n",
    "    else:\n",
    "        fsize = (16, 4 * (plotrows + 1))\n",
    "        plt.figure(figsize=fsize, dpi=dpi)\n",
    "    \n",
    "    tce, tces_file = get_tce(tic, item_no)\n",
    "    period = tce.Period.values.item()\n",
    "    epoc = tce.Epoc.values.item()\n",
    "    duration = tce.Duration.values.item()\n",
    "\n",
    "    print(f'Epoc: {epoc}\\nPeriod: {period}\\nDuration: {duration}')\n",
    "    \n",
    "    def config_abs_plot(title):\n",
    "        plt.legend()\n",
    "        plt.title(title)\n",
    "        if abs_xlim:\n",
    "            if abs_xlim == '3p':\n",
    "                minx = min(td) + abs_offset * period\n",
    "                maxx = minx + 3.5 * period\n",
    "                plt.xlim(minx, maxx)\n",
    "            else:\n",
    "                plt.xlim(*abs_xlim)\n",
    "        if abs_ylim:\n",
    "            if abs_ylim == '3%':\n",
    "                miny = np.percentile(fs[~np.isnan(fs)], 3)\n",
    "                maxy = np.percentile(fs[~np.isnan(fs)], 97)\n",
    "                plt.ylim(miny, maxy)\n",
    "            else:\n",
    "                plt.ylim(*abs_ylim)\n",
    "                \n",
    "    nplotted = 0\n",
    "    def splt(c=1):\n",
    "        nonlocal nplotted\n",
    "        if large:\n",
    "            plt.figure(figsize=fsize, dpi=dpi)\n",
    "        else:\n",
    "            plt.subplot(plotrows, plotcols // c, (nplotted // c) + 1)\n",
    "        nplotted += c\n",
    "    \n",
    "    \n",
    "    file_names = tess_io.tess_filenames(tic, tess_data_dir)\n",
    "    f = fits.open(file_names)\n",
    "\n",
    "    td = f[1].data[\"TIME\"]\n",
    "    fd = f[1].data[\"KSPSAP_FLUX\"]\n",
    "    fs = f[1].data[\"SAP_FLUX\"]\n",
    "    \n",
    "#     splt()\n",
    "#     plt.plot(td, fs, '-', alpha=0.6, color=g, label='SAP')\n",
    "#     config_abs_plot('fits data')\n",
    "    \n",
    "    \n",
    "    ut, uf = preprocess.read_and_process_light_curve(tic, tess_data_dir, 'SAP_FLUX')\n",
    "\n",
    "    input_mask = preprocess.get_spline_mask(ut, period, epoc, duration)\n",
    "    sf, mdata = keplersplinev2.choosekeplersplinev2(\n",
    "        ut, uf, input_mask=input_mask, return_metadata=True,\n",
    "        fixed_bkspace=bkspace,\n",
    "    )\n",
    "\n",
    "    splt()\n",
    "    plt.plot(ut, uf, '-', alpha=0.6, color=g, label='SAP')\n",
    "    plt.plot(ut[input_mask], sf[input_mask], 'x', markersize=3, color=k, label='spline (out of transit)')\n",
    "    plt.plot(ut[~input_mask], sf[~input_mask], 'o', markersize=3, color=o, label='spline (transit)')\n",
    "    config_abs_plot(f'raw | bkspace: {mdata.bkspace}')\n",
    "    \n",
    "    \n",
    "    ut, nf, fm = preprocess.detrend_and_filter(\n",
    "        tic, ut, uf, period, epoc, duration, bkspace)\n",
    "    sft, sff, sfn, sftm = preprocess.phase_fold_and_sort_light_curve(\n",
    "        ut, nf, input_mask, period, epoc)\n",
    "\n",
    "#     splt()\n",
    "#     plt.plot(sft, sff, 'o', markersize=3, alpha=0.6, c=o, label='spline')\n",
    "#     sff_filtered = np.where((sff > 1.5) | (sff < -0.5), 0, sff)\n",
    "#     if len(sff_filtered):\n",
    "#         mask = np.where(sftm, 1, min(sff_filtered))\n",
    "#         plt.plot(sft, mask, '-', markersize=1, alpha=0.6, c=r, label='OOT')\n",
    "#         title = f'phase folded ({int(max(sfn) + 1)} folds) / filtered'\n",
    "#     else:\n",
    "#         if len(sfn):\n",
    "#             title = f'phase folded ({int(max(sfn) + 1)} folds) | WARNING: filtering removed all data'\n",
    "#         else:\n",
    "#             title = f'phase folded | WARNING: filtering removed all data'\n",
    "#     plt.legend(prop={'size': 14})\n",
    "#     plt.title(title)\n",
    "\n",
    "\n",
    "    generate_input_records.FLAGS = generate_input_records.parser.parse_args([\n",
    "      '--tess_data_dir', '/mnt/tess/lc',\n",
    "      '--output_dir', '/dev/null',\n",
    "      '--input_tce_csv_file', tces_file,\n",
    "    ])\n",
    "\n",
    "\n",
    "    row = list(tce.iterrows())[0][1]\n",
    "\n",
    "    ex = generate_input_records._process_tce(row, bkspace)\n",
    "    \n",
    "    \n",
    "    def plot_w_scatter(view, std, title, mask=None):\n",
    "        splt()\n",
    "        view = np.array(ex.features.feature[view].float_list.value)\n",
    "        std = np.array(ex.features.feature[std].float_list.value)\n",
    "        plt.plot(std, color=o, label='Standard deviation')\n",
    "        plt.plot(view, color=n, label='Filtered mean')\n",
    "        if mask:\n",
    "            plt.plot(\n",
    "                -np.array(ex.features.feature[mask].float_list.value), color=r, alpha=0.5, linestyle='--',\n",
    "            label='Transit mask')\n",
    "        plt.legend(prop={'size': 14})\n",
    "        plt.title(title)\n",
    "        \n",
    "    def plot_norm(view, title):\n",
    "        splt()\n",
    "        plt.plot(ex.features.feature[view].float_list.value, color=n)\n",
    "        plt.ylim(-1, 1)\n",
    "        plt.title(title)    \n",
    "        \n",
    "        \n",
    "    def plot_segments(view, tag):\n",
    "        splt(c=3)\n",
    "        img = ex.features.feature[view].float_list.value\n",
    "        img = np.reshape(img, (-1, 14))\n",
    "        n_transits = img.shape[1] // 2\n",
    "        for i in range(n_transits):\n",
    "            view = img[:, 2 * i]\n",
    "            mask = img[:, 2 * i + 1] > 0\n",
    "            plt.plot(np.where(mask, view, np.nan), marker='.', label=f'Sample transit {i}')\n",
    "        plt.legend(prop={'size': 14})\n",
    "        plt.title(f'{n_transits} sample segments, densest first, ties broken at random {tag}')\n",
    "        \n",
    "#     plot_w_scatter('global_view', 'global_std', 'global')\n",
    "#     plot_w_scatter('global_view_0.3', 'global_std_0.3', 'global @ 0.3')\n",
    "#     plot_w_scatter('global_view_5.0', 'global_std_5.0', 'global @ 5.0')\n",
    "\n",
    "    sec_phase = ex.features.feature['secondary_phase'].float_list.value[0]\n",
    "    loc_scale = ex.features.feature['local_scale'].float_list.value[0]\n",
    "    sec_scale = ex.features.feature['secondary_scale'].float_list.value[0]\n",
    "#     plot_w_scatter('global_view', 'global_std', 'global', 'global_transit_mask')\n",
    "#     plot_w_scatter('local_view', 'local_std', f'local / {loc_scale:0.3}')\n",
    "#     plot_w_scatter('secondary_view', 'secondary_std', f'secondary ({sec_phase:0.2}) / {sec_scale:0.5}')\n",
    "#     plot_norm('global_view_half_period', 'global half period')\n",
    "#     plot_norm('global_view_double_period', 'global double period (shifted)')\n",
    "#     splt()\n",
    "#     plot_segments('sample_segments_view', '')\n",
    "\n",
    "#     sec_phase = ex.features.feature['secondary_phase_0.3'].float_list.value[0]\n",
    "#     loc_scale = ex.features.feature['local_scale_0.3'].float_list.value[0]\n",
    "#     sec_scale = ex.features.feature['secondary_scale_0.3'].float_list.value[0]\n",
    "#     plot_w_scatter('global_view_0.3', 'global_std_0.3', 'global @ 0.3', 'global_transit_mask_0.3')\n",
    "#     plot_w_scatter('local_view_0.3', 'local_std_0.3', f'local @ 0.3 / {loc_scale:0.3}')\n",
    "#     plot_w_scatter('secondary_view_0.3', 'secondary_std_0.3', f'secondary @ 0.3 ({sec_phase:0.2}) / {sec_scale:0.5}')\n",
    "#     plot_norm('global_view_half_period_0.3', 'global half period @ 0.3')\n",
    "#     plot_norm('global_view_double_period_0.3', 'global double period (shifted) @ 0.3')\n",
    "#     splt()\n",
    "#     plot_segments('sample_segments_view_0.3', ' @ 0.3')\n",
    "\n",
    "#     sec_phase = ex.features.feature['secondary_phase_5.0'].float_list.value[0]\n",
    "#     loc_scale = ex.features.feature['local_scale_5.0'].float_list.value[0]\n",
    "#     sec_scale = ex.features.feature['secondary_scale_5.0'].float_list.value[0]\n",
    "#     plot_w_scatter('global_view_5.0', 'global_std_5.0', 'global @ 5.0', 'global_transit_mask_5.0')\n",
    "#     plot_w_scatter('local_view_5.0', 'local_std_5.0', f'local @ 5.0 / {loc_scale:0.3}')\n",
    "#     plot_w_scatter('secondary_view_5.0', 'secondary_std_5.0', f'secondary @ 5.0 ({sec_phase:0.2}) / {sec_scale:0.5}')\n",
    "#     plot_norm('global_view_half_period_5.0', 'global half period @ 5.0')\n",
    "#     plot_norm('global_view_double_period_5.0', 'global double period (shifted) @ 5.0')\n",
    "#     splt()\n",
    "#     plot_segments('sample_segments_view_5.0', ' @ 5.0')\n",
    "    \n",
    "#     try:\n",
    "#         im = plt.imread(f'{reports_dir}/TIC{tic}.png')\n",
    "#         plt.figure(figsize=(12, 8), dpi=dpi)\n",
    "#         plt.axis('off')\n",
    "#         _ = plt.imshow(im)\n",
    "#     except FileNotFoundError:\n",
    "#         print('-- no report file --')\n",
    "\n",
    "plot_all(354444731, large=True)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_all(tic,\n",
    "             large=False,\n",
    "             dpi=100,\n",
    "             bkspace=None,\n",
    "             abs_xlim=None,\n",
    "             abs_offset=0,\n",
    "             abs_ylim=None,\n",
    "             rpt_only=False,\n",
    "             item_no=0):\n",
    "    tess_data_dir = f'/mnt/tess/lc'\n",
    "    reports_dir = f'/mnt/tess/rpt/png'\n",
    "\n",
    "    if rpt_only:\n",
    "        plt.figure(figsize=(12, 8), dpi=dpi)\n",
    "        plt.axis('off')\n",
    "        im = plt.imread(f'{reports_dir}/TIC{tic}.png')\n",
    "        _ = plt.imshow(im)\n",
    "        plt.show()\n",
    "        print(get_tce(tic, item_no)[0])\n",
    "        return\n",
    "        \n",
    "    b = plt.get_cmap('tab20')(0)\n",
    "    b2 = plt.get_cmap('tab20')(1)\n",
    "    o = plt.get_cmap('tab20')(2)\n",
    "    o2 = plt.get_cmap('tab20')(3)\n",
    "    g = plt.get_cmap('tab20')(4)\n",
    "    r = plt.get_cmap('tab20')(6)\n",
    "    n = plt.get_cmap('tab20')(8)\n",
    "    k = plt.get_cmap('tab20')(14)\n",
    "\n",
    "    plotrows = 1\n",
    "    plotcols = 3\n",
    "    \n",
    "    if large:\n",
    "        fsize = (5, 3)\n",
    "    else:\n",
    "        fsize = (16, 4 * (plotrows + 1))\n",
    "        plt.figure(figsize=fsize, dpi=dpi)\n",
    "    \n",
    "    tce, tces_file = get_tce(tic, item_no)\n",
    "    period = tce.Period.values.item()\n",
    "    epoc = tce.Epoc.values.item()\n",
    "    duration = tce.Duration.values.item()\n",
    "\n",
    "    print(f'Epoc: {epoc}\\nPeriod: {period}\\nDuration: {duration}')\n",
    "    \n",
    "    def config_abs_plot(title):\n",
    "#         plt.legend()\n",
    "        plt.title(title)\n",
    "        if abs_xlim:\n",
    "            if abs_xlim == '3p':\n",
    "                minx = min(td) + abs_offset * period\n",
    "                maxx = minx + 3.5 * period\n",
    "                plt.xlim(minx, maxx)\n",
    "            else:\n",
    "                plt.xlim(*abs_xlim)\n",
    "        if abs_ylim:\n",
    "            if abs_ylim == '3%':\n",
    "                miny = np.percentile(fs[~np.isnan(fs)], 3)\n",
    "                maxy = np.percentile(fs[~np.isnan(fs)], 97)\n",
    "                plt.ylim(miny, maxy)\n",
    "            else:\n",
    "                plt.ylim(*abs_ylim)\n",
    "                \n",
    "    nplotted = 0\n",
    "    def splt(c=1):\n",
    "        nonlocal nplotted\n",
    "        if large:\n",
    "            plt.figure(figsize=fsize, dpi=dpi)\n",
    "        else:\n",
    "            plt.subplot(plotrows, plotcols // c, (nplotted // c) + 1)\n",
    "        nplotted += c\n",
    "    \n",
    "    \n",
    "    file_names = tess_io.tess_filenames(tic, tess_data_dir)\n",
    "    f = fits.open(file_names)\n",
    "\n",
    "    td = f[1].data[\"TIME\"]\n",
    "    fd = f[1].data[\"KSPSAP_FLUX\"]\n",
    "    fs = f[1].data[\"SAP_FLUX\"]\n",
    "    \n",
    "#     splt()\n",
    "#     plt.plot(td, fs, '-', alpha=0.6, color=g, label='SAP')\n",
    "#     config_abs_plot('fits data')\n",
    "    \n",
    "    \n",
    "    ut, uf = preprocess.read_and_process_light_curve(tic, tess_data_dir, 'SAP_FLUX')\n",
    "\n",
    "    input_mask = preprocess.get_spline_mask(ut, period, epoc, duration)\n",
    "    sf, mdata = keplersplinev2.choosekeplersplinev2(\n",
    "        ut, uf, input_mask=input_mask, return_metadata=True,\n",
    "        fixed_bkspace=bkspace,\n",
    "    )\n",
    "\n",
    "    splt()\n",
    "    plt.plot(ut, uf, '-', alpha=0.6, color=g, label='SAP')\n",
    "    plt.plot(ut[input_mask], sf[input_mask], 'x', markersize=3, color=k, label='spline (out of transit)')\n",
    "    plt.plot(ut[~input_mask], sf[~input_mask], 'o', markersize=3, color=o, label='spline (transit)')\n",
    "    config_abs_plot(f'raw | bkspace: {mdata.bkspace}')\n",
    "    \n",
    "    \n",
    "    ut, nf, fm = preprocess.detrend_and_filter(\n",
    "        tic, ut, uf, period, epoc, duration, bkspace)\n",
    "    sft, sff, sfn, sftm = preprocess.phase_fold_and_sort_light_curve(\n",
    "        ut, nf, input_mask, period, epoc)\n",
    "\n",
    "#     splt()\n",
    "#     plt.plot(sft, sff, 'o', markersize=3, alpha=0.6, c=o, label='spline')\n",
    "#     sff_filtered = np.where((sff > 1.5) | (sff < -0.5), 0, sff)\n",
    "#     if len(sff_filtered):\n",
    "#         mask = np.where(sftm, 1, min(sff_filtered))\n",
    "#         plt.plot(sft, mask, '-', markersize=1, alpha=0.6, c=r, label='OOT')\n",
    "#         title = f'phase folded ({int(max(sfn) + 1)} folds) / filtered'\n",
    "#     else:\n",
    "#         if len(sfn):\n",
    "#             title = f'phase folded ({int(max(sfn) + 1)} folds) | WARNING: filtering removed all data'\n",
    "#         else:\n",
    "#             title = f'phase folded | WARNING: filtering removed all data'\n",
    "#     plt.legend(prop={'size': 14})\n",
    "#     plt.title(title)\n",
    "\n",
    "\n",
    "    generate_input_records.FLAGS = generate_input_records.parser.parse_args([\n",
    "      '--tess_data_dir', '/mnt/tess/lc',\n",
    "      '--output_dir', '/dev/null',\n",
    "      '--input_tce_csv_file', tces_file,\n",
    "    ])\n",
    "\n",
    "\n",
    "    row = list(tce.iterrows())[0][1]\n",
    "\n",
    "    ex = generate_input_records._process_tce(row, bkspace)\n",
    "    \n",
    "    \n",
    "    def plot_w_scatter(view, std, title, mask=None):\n",
    "        splt()\n",
    "        view = np.array(ex.features.feature[view].float_list.value)\n",
    "        std = np.array(ex.features.feature[std].float_list.value)\n",
    "        plt.plot(std, color=o, label='Standard deviation')\n",
    "        plt.plot(view, color=n, label='Filtered mean')\n",
    "        if mask:\n",
    "            plt.plot(\n",
    "                -np.array(ex.features.feature[mask].float_list.value), color=r, alpha=0.5, linestyle='--',\n",
    "            label='Transit mask')\n",
    "#         plt.legend(prop={'size': 14})\n",
    "        plt.title(title)\n",
    "        \n",
    "    def plot_norm(view, title):\n",
    "        splt()\n",
    "        plt.plot(ex.features.feature[view].float_list.value, color=n)\n",
    "        plt.ylim(-1, 1)\n",
    "        plt.title(title)    \n",
    "        \n",
    "        \n",
    "    def plot_segments(view, tag):\n",
    "        splt(c=3)\n",
    "        img = ex.features.feature[view].float_list.value\n",
    "        img = np.reshape(img, (-1, 14))\n",
    "        n_transits = img.shape[1] // 2\n",
    "        for i in range(n_transits):\n",
    "            view = img[:, 2 * i]\n",
    "            mask = img[:, 2 * i + 1] > 0\n",
    "            plt.plot(np.where(mask, view, np.nan), marker='.', label=f'Sample transit {i}')\n",
    "        plt.legend(prop={'size': 14})\n",
    "        plt.title(f'{n_transits} sample segments, densest first, ties broken at random {tag}')\n",
    "        \n",
    "#     plot_w_scatter('global_view', 'global_std', 'global')\n",
    "#     plot_w_scatter('global_view_0.3', 'global_std_0.3', 'global @ 0.3')\n",
    "#     plot_w_scatter('global_view_5.0', 'global_std_5.0', 'global @ 5.0')\n",
    "\n",
    "    sec_phase = ex.features.feature['secondary_phase'].float_list.value[0]\n",
    "    loc_scale = ex.features.feature['local_scale'].float_list.value[0]\n",
    "    sec_scale = ex.features.feature['secondary_scale'].float_list.value[0]\n",
    "#     plot_w_scatter('global_view', 'global_std', 'global', 'global_transit_mask')\n",
    "#     plot_w_scatter('local_view', 'local_std', f'local / {loc_scale:0.3}')\n",
    "#     plot_w_scatter('secondary_view', 'secondary_std', f'secondary ({sec_phase:0.2}) / {sec_scale:0.5}')\n",
    "#     plot_norm('global_view_half_period', 'global half period')\n",
    "#     plot_norm('global_view_double_period', 'global double period (shifted)')\n",
    "#     splt()\n",
    "#     plot_segments('sample_segments_view', '')\n",
    "\n",
    "#     sec_phase = ex.features.feature['secondary_phase_0.3'].float_list.value[0]\n",
    "#     loc_scale = ex.features.feature['local_scale_0.3'].float_list.value[0]\n",
    "#     sec_scale = ex.features.feature['secondary_scale_0.3'].float_list.value[0]\n",
    "    plot_w_scatter('global_view_0.3', 'global_std_0.3', 'global @ 0.3', 'global_transit_mask_0.3')\n",
    "#     plot_w_scatter('local_view_0.3', 'local_std_0.3', f'local @ 0.3 / {loc_scale:0.3}')\n",
    "#     plot_w_scatter('secondary_view_0.3', 'secondary_std_0.3', f'secondary @ 0.3 ({sec_phase:0.2}) / {sec_scale:0.5}')\n",
    "#     plot_norm('global_view_half_period_0.3', 'global half period @ 0.3')\n",
    "#     plot_norm('global_view_double_period_0.3', 'global double period (shifted) @ 0.3')\n",
    "#     splt()\n",
    "#     plot_segments('sample_segments_view_0.3', ' @ 0.3')\n",
    "\n",
    "#     sec_phase = ex.features.feature['secondary_phase_5.0'].float_list.value[0]\n",
    "#     loc_scale = ex.features.feature['local_scale_5.0'].float_list.value[0]\n",
    "#     sec_scale = ex.features.feature['secondary_scale_5.0'].float_list.value[0]\n",
    "#     plot_w_scatter('global_view_5.0', 'global_std_5.0', 'global @ 5.0', 'global_transit_mask_5.0')\n",
    "#     plot_w_scatter('local_view_5.0', 'local_std_5.0', f'local @ 5.0 / {loc_scale:0.3}')\n",
    "#     plot_w_scatter('secondary_view_5.0', 'secondary_std_5.0', f'secondary @ 5.0 ({sec_phase:0.2}) / {sec_scale:0.5}')\n",
    "#     plot_norm('global_view_half_period_5.0', 'global half period @ 5.0')\n",
    "#     plot_norm('global_view_double_period_5.0', 'global double period (shifted) @ 5.0')\n",
    "#     splt()\n",
    "#     plot_segments('sample_segments_view_5.0', ' @ 5.0')\n",
    "    \n",
    "#     try:\n",
    "#         im = plt.imread(f'{reports_dir}/TIC{tic}.png')\n",
    "#         plt.figure(figsize=(12, 8), dpi=dpi)\n",
    "#         plt.axis('off')\n",
    "#         _ = plt.imshow(im)\n",
    "#     except FileNotFoundError:\n",
    "#         print('-- no report file --')\n",
    "\n",
    "plot_all(354444731, large=True, bkspace=0.3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "from astropy.io import fits\n",
    "from astronet.preprocess import generate_input_records\n",
    "from astronet.preprocess import preprocess\n",
    "from light_curve_util import keplersplinev2\n",
    "from light_curve_util import tess_io\n",
    "from matplotlib import pyplot as plt\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "\n",
    "\n",
    "tce_tables = {}\n",
    "\n",
    "tce_files = [\n",
    "    '/mnt/tess/astronet/tces-toi.csv',\n",
    "#     '/mnt/tess/astronet/tces-toi-bls-vs-tev.csv',\n",
    "    '/mnt/tess/astronet/tces-s33_cam1_sample.csv',\n",
    "    '/mnt/tess/astronet/tces-s33_cam2ccd14_sample.csv',\n",
    "    '/mnt/tess/astronet/tces-v6-train.csv',\n",
    "    '/mnt/tess/astronet/tces-v6-val.csv',\n",
    "    '/mnt/tess/astronet/tces-v6-test.csv',\n",
    "]\n",
    "\n",
    "\n",
    "def get_tce(tic, item_no):\n",
    "    for f in tce_files:\n",
    "        if f not in tce_tables:\n",
    "            tce_tables[f] = pd.read_csv(f, header=0)\n",
    "        tce_table = tce_tables[f]\n",
    "        tce = tce_table[tce_table.tic_id == tic]\n",
    "        if not len(tce):\n",
    "            continue\n",
    "        if len(tce) > 1:\n",
    "            tce = tce[tce.index == tce.index.values[item_no]]\n",
    "            if 'Source' in tce:\n",
    "                print('Source:', 'BLS' if tce.Source.values.item() == 2 else 'TEV')\n",
    "        return tce, f\n",
    "    raise ValueError(f'no TCE data for {tic}')\n",
    "\n",
    "\n",
    "def plot_all(tic,\n",
    "             large=False,\n",
    "             dpi=100,\n",
    "             bkspace=None,\n",
    "             abs_xlim=None,\n",
    "             abs_offset=0,\n",
    "             abs_ylim=None,\n",
    "             rpt_only=False,\n",
    "             item_no=0):\n",
    "    tess_data_dir = f'/mnt/tess/lc'\n",
    "    reports_dir = f'/mnt/tess/rpt/png'\n",
    "\n",
    "    if rpt_only:\n",
    "        plt.figure(figsize=(12, 8), dpi=dpi)\n",
    "        plt.axis('off')\n",
    "        im = plt.imread(f'{reports_dir}/TIC{tic}.png')\n",
    "        _ = plt.imshow(im)\n",
    "        plt.show()\n",
    "        print(get_tce(tic, item_no)[0])\n",
    "        return\n",
    "        \n",
    "    b = plt.get_cmap('tab20')(0)\n",
    "    b2 = plt.get_cmap('tab20')(1)\n",
    "    o = plt.get_cmap('tab20')(2)\n",
    "    o2 = plt.get_cmap('tab20')(3)\n",
    "    g = plt.get_cmap('tab20')(4)\n",
    "    r = plt.get_cmap('tab20')(6)\n",
    "    n = plt.get_cmap('tab20')(8)\n",
    "    k = plt.get_cmap('tab20')(14)\n",
    "\n",
    "    plotrows = 11\n",
    "    plotcols = 3\n",
    "    \n",
    "    if large:\n",
    "        fsize = (16, 9)\n",
    "    else:\n",
    "        fsize = (16, 3 * (plotrows + 1))\n",
    "        plt.figure(figsize=fsize, dpi=dpi)\n",
    "    \n",
    "    tce, tces_file = get_tce(tic, item_no)\n",
    "    period = tce.Period.values.item()\n",
    "    epoc = tce.Epoc.values.item()\n",
    "    duration = tce.Duration.values.item()\n",
    "\n",
    "    print(f'Epoc: {epoc}\\nPeriod: {period}\\nDuration: {duration}')\n",
    "    \n",
    "    def config_abs_plot(title):\n",
    "        plt.legend()\n",
    "        plt.title(title)\n",
    "        if abs_xlim:\n",
    "            if abs_xlim == '3p':\n",
    "                minx = min(td) + abs_offset * period\n",
    "                maxx = minx + 3.5 * period\n",
    "                plt.xlim(minx, maxx)\n",
    "            else:\n",
    "                plt.xlim(*abs_xlim)\n",
    "        if abs_ylim:\n",
    "            if abs_ylim == '3%':\n",
    "                miny = np.percentile(fs[~np.isnan(fs)], 3)\n",
    "                maxy = np.percentile(fs[~np.isnan(fs)], 97)\n",
    "                plt.ylim(miny, maxy)\n",
    "            else:\n",
    "                plt.ylim(*abs_ylim)\n",
    "                \n",
    "    nplotted = 0\n",
    "    def splt(c=1):\n",
    "        nonlocal nplotted\n",
    "        if large:\n",
    "            plt.figure(figsize=fsize, dpi=dpi)\n",
    "        else:\n",
    "            plt.subplot(plotrows, plotcols // c, (nplotted // c) + 1)\n",
    "        nplotted += c\n",
    "    \n",
    "    \n",
    "    file_names = tess_io.tess_filenames(tic, tess_data_dir)\n",
    "    f = fits.open(file_names)\n",
    "\n",
    "    td = f[1].data[\"TIME\"]\n",
    "    fd = f[1].data[\"KSPSAP_FLUX\"]\n",
    "    fs = f[1].data[\"SAP_FLUX\"]\n",
    "    \n",
    "#     splt()\n",
    "#     plt.plot(td, fs, '-', alpha=0.6, color=g, label='SAP')\n",
    "#     config_abs_plot('fits data')\n",
    "    \n",
    "    \n",
    "    splt()\n",
    "    ut, uf = preprocess.read_and_process_light_curve(tic, tess_data_dir, 'SAP_FLUX')\n",
    "\n",
    "    input_mask = preprocess.get_spline_mask(ut, period, epoc, duration)\n",
    "    sf, mdata = keplersplinev2.choosekeplersplinev2(\n",
    "        ut, uf, input_mask=input_mask, return_metadata=True,\n",
    "        fixed_bkspace=bkspace,\n",
    "    )\n",
    "\n",
    "    plt.plot(ut, uf, '-', alpha=0.6, color=g, label='SAP')\n",
    "    plt.plot(ut[input_mask], sf[input_mask], 'x', markersize=3, color=k, label='spline (OOT)')\n",
    "    plt.plot(ut[~input_mask], sf[~input_mask], 'o', markersize=3, color=o, label='spline')\n",
    "    config_abs_plot(f'raw | bkspace: {mdata.bkspace}')\n",
    "    \n",
    "    \n",
    "    ut, nf, fm = preprocess.detrend_and_filter(\n",
    "        tic, ut, uf, period, epoc, duration, bkspace)\n",
    "    sft, sff, sfn, sftm = preprocess.phase_fold_and_sort_light_curve(\n",
    "        ut, nf, input_mask, period, epoc)\n",
    "\n",
    "    splt()\n",
    "    plt.plot(sft, sff, 'o', markersize=3, alpha=0.6, c=o, label='spline')\n",
    "    sff_filtered = np.where((sff > 1.5) | (sff < -0.5), 0, sff)\n",
    "    if len(sff_filtered):\n",
    "        mask = np.where(sftm, 1, min(sff_filtered))\n",
    "        plt.plot(sft, mask, '-', markersize=1, alpha=0.6, c=r, label='OOT')\n",
    "        title = f'phase folded ({int(max(sfn) + 1)} folds) / filtered'\n",
    "    else:\n",
    "        if len(sfn):\n",
    "            title = f'phase folded ({int(max(sfn) + 1)} folds) | WARNING: filtering removed all data'\n",
    "        else:\n",
    "            title = f'phase folded | WARNING: filtering removed all data'\n",
    "    plt.legend()\n",
    "    plt.title(title)\n",
    "\n",
    "\n",
    "    generate_input_records.FLAGS = generate_input_records.parser.parse_args([\n",
    "      '--tess_data_dir', '/mnt/tess/lc',\n",
    "      '--output_dir', '/dev/null',\n",
    "      '--input_tce_csv_file', tces_file,\n",
    "    ])\n",
    "\n",
    "\n",
    "    row = list(tce.iterrows())[0][1]\n",
    "\n",
    "    ex = generate_input_records._process_tce(row, bkspace)\n",
    "    \n",
    "    \n",
    "    def plot_w_scatter(view, std, title, mask=None):\n",
    "        splt()\n",
    "        view = np.array(ex.features.feature[view].float_list.value)\n",
    "        std = np.array(ex.features.feature[std].float_list.value)\n",
    "        plt.plot(std, color=o)\n",
    "        plt.plot(view, color=n)\n",
    "        if mask:\n",
    "            plt.plot(\n",
    "                -np.array(ex.features.feature[mask].float_list.value), color=r, alpha=0.5, linestyle='--')\n",
    "        plt.title(title)\n",
    "        \n",
    "    def plot_norm(view, title):\n",
    "        splt()\n",
    "        plt.plot(ex.features.feature[view].float_list.value, color=n)\n",
    "        plt.ylim(-1, 1)\n",
    "        plt.title(title)    \n",
    "        \n",
    "        \n",
    "    def plot_segments(view, tag):\n",
    "        splt(c=3)\n",
    "        img = ex.features.feature[view].float_list.value\n",
    "        img = np.reshape(img, (-1, 14))\n",
    "        n_transits = img.shape[1] // 2\n",
    "        for i in range(n_transits):\n",
    "            view = img[:, 2 * i]\n",
    "            mask = img[:, 2 * i + 1] > 0\n",
    "            plt.plot(np.where(mask, view, np.nan), marker='.')\n",
    "        plt.title(f'{n_transits} sample segments, densest first, ties broken at random {tag}')\n",
    "        \n",
    "#     plot_w_scatter('global_view', 'global_std', 'global')\n",
    "#     plot_w_scatter('global_view_0.3', 'global_std_0.3', 'global @ 0.3')\n",
    "#     plot_w_scatter('global_view_5.0', 'global_std_5.0', 'global @ 5.0')\n",
    "\n",
    "    sec_phase = ex.features.feature['secondary_phase'].float_list.value[0]\n",
    "    loc_scale = ex.features.feature['local_scale'].float_list.value[0]\n",
    "    sec_scale = ex.features.feature['secondary_scale'].float_list.value[0]\n",
    "    plot_w_scatter('global_view', 'global_std', 'global', 'global_transit_mask')\n",
    "    plot_w_scatter('local_view', 'local_std', f'local / {loc_scale:0.3}')\n",
    "    plot_w_scatter('secondary_view', 'secondary_std', f'secondary ({sec_phase:0.2}) / {sec_scale:0.5}')\n",
    "#     plot_norm('global_view_half_period', 'global half period')\n",
    "    plot_norm('global_view_double_period', 'global double period (shifted)')\n",
    "    splt()\n",
    "    plot_segments('sample_segments_view', '')\n",
    "\n",
    "#     sec_phase = ex.features.feature['secondary_phase_0.3'].float_list.value[0]\n",
    "#     loc_scale = ex.features.feature['local_scale_0.3'].float_list.value[0]\n",
    "#     sec_scale = ex.features.feature['secondary_scale_0.3'].float_list.value[0]\n",
    "#     plot_w_scatter('global_view_0.3', 'global_std_0.3', 'global @ 0.3', 'global_transit_mask_0.3')\n",
    "#     plot_w_scatter('local_view_0.3', 'local_std_0.3', f'local @ 0.3 / {loc_scale:0.3}')\n",
    "#     plot_w_scatter('secondary_view_0.3', 'secondary_std_0.3', f'secondary @ 0.3 ({sec_phase:0.2}) / {sec_scale:0.5}')\n",
    "#     plot_norm('global_view_half_period_0.3', 'global half period @ 0.3')\n",
    "#     plot_norm('global_view_double_period_0.3', 'global double period (shifted) @ 0.3')\n",
    "#     splt()\n",
    "#     plot_segments('sample_segments_view_0.3', ' @ 0.3')\n",
    "\n",
    "#     sec_phase = ex.features.feature['secondary_phase_5.0'].float_list.value[0]\n",
    "#     loc_scale = ex.features.feature['local_scale_5.0'].float_list.value[0]\n",
    "#     sec_scale = ex.features.feature['secondary_scale_5.0'].float_list.value[0]\n",
    "#     plot_w_scatter('global_view_5.0', 'global_std_5.0', 'global @ 5.0', 'global_transit_mask_5.0')\n",
    "#     plot_w_scatter('local_view_5.0', 'local_std_5.0', f'local @ 5.0 / {loc_scale:0.3}')\n",
    "#     plot_w_scatter('secondary_view_5.0', 'secondary_std_5.0', f'secondary @ 5.0 ({sec_phase:0.2}) / {sec_scale:0.5}')\n",
    "#     plot_norm('global_view_half_period_5.0', 'global half period @ 5.0')\n",
    "#     plot_norm('global_view_double_period_5.0', 'global double period (shifted) @ 5.0')\n",
    "#     splt()\n",
    "#     plot_segments('sample_segments_view_5.0', ' @ 5.0')\n",
    "    \n",
    "#     try:\n",
    "#         im = plt.imread(f'{reports_dir}/TIC{tic}.png')\n",
    "#         plt.figure(figsize=(12, 8), dpi=dpi)\n",
    "#         plt.axis('off')\n",
    "#         _ = plt.imshow(im)\n",
    "#     except FileNotFoundError:\n",
    "#         print('-- no report file --')\n",
    "\n",
    "plot_all(160440924,\n",
    "#          rpt_only=True,\n",
    "#          large=True,\n",
    "#          bkspace=5.0,\n",
    "#          abs_xlim='3p',\n",
    "#          abs_offset=6,\n",
    "#          abs_xlim=(1300, 2500),\n",
    "#          abs_ylim='3%',\n",
    "#          abs_ylim=(0.975, 1.01),\n",
    "#          item_no=10\n",
    "        )"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
